"""
SQL Server 数据库管理类

该模块提供了 SQL Server 数据库的连接管理和常用操作功能，包括查询执行、数据插入、更新、删除等。
"""
import pyodbc
from typing import Dict, List, Tuple, Optional, Any
from modules.logger import LogManager


class SQLServerManager:
    """
    SQL Server 数据库管理类，封装了常用的数据库操作方法
    """
    def __init__(self,
                 server: str = 'localhost',
                 database: str = 'master',
                 username: str = 'sa',
                 password: str = '',
                 port: int = 1433,
                 driver: str = '{ODBC Driver 17 for SQL Server}',
                 logger_name: str = 'sqlserver_manager'):
        """
        初始化 SQL Server 连接参数

        Args:
            server: 服务器地址
            database: 数据库名称
            username: 登录用户名
            password: 登录密码
            port: 数据库端口
            driver: ODBC 驱动名称
            logger_name: 日志记录器名称
        """
        self.server = server
        self.database = database
        self.username = username
        self.password = password
        self.port = port
        self.driver = driver
        self.connection: Optional[pyodbc.Connection] = None
        self.cursor: Optional[pyodbc.Cursor] = None
        
        # 初始化日志
        self.logger = LogManager.get_both_logger(name=logger_name)

    def connect(self) -> bool:
        """
        建立数据库连接

        Returns:
            bool: 连接成功返回 True，失败返回 False
        """
        try:
            connection_string = (
                f"DRIVER={self.driver};"
                f"SERVER={self.server},{self.port};"
                f"DATABASE={self.database};"
                f"UID={self.username};"
                f"PWD={self.password};"
                "TrustServerCertificate=yes"
            )
            
            self.connection = pyodbc.connect(connection_string)
            self.cursor = self.connection.cursor()
            self.logger.info(f"成功连接到 SQL Server: {self.server}/{self.database}")
            return True
        except Exception as e:
            self.logger.error(f"数据库连接失败: {str(e)}", exc_info=True)
            self.connection = None
            self.cursor = None
            return False

    def close(self) -> None:
        """
        关闭数据库连接
        """
        if self.cursor:
            try:
                self.cursor.close()
            except Exception as e:
                self.logger.warning(f"关闭游标时出错: {str(e)}")
            finally:
                self.cursor = None

        if self.connection:
            try:
                self.connection.close()
                self.logger.info(f"已关闭与 {self.server}/{self.database} 的连接")
            except Exception as e:
                self.logger.warning(f"关闭连接时出错: {str(e)}")
            finally:
                self.connection = None

    def _ensure_connection(self) -> bool:
        """
        确保数据库连接处于打开状态

        Returns:
            bool: 连接有效返回 True，否则返回 False
        """
        if not self.connection or self.connection.state == 0:
            return self.connect()
        return True

    def execute_query(self, query: str, params: Optional[Tuple[Any, ...]] = None) -> List[Dict[str, Any]]:
        """
        执行查询语句并返回结果

        Args:
            query: SQL 查询语句
            params: 查询参数（用于参数化查询）

        Returns:
            List[Dict[str, Any]]: 查询结果列表，每个元素为一行数据的字典
        """
        if not self._ensure_connection():
            self.logger.error("无法执行查询，数据库连接失败")
            return []

        try:
            if params:
                self.cursor.execute(query, params)
            else:
                self.cursor.execute(query)

            # 获取列名
            columns = [column[0] for column in self.cursor.description]
            # 获取结果并转换为字典列表
            results = [dict(zip(columns, row)) for row in self.cursor.fetchall()]
            
            self.logger.info(f"查询执行成功，返回 {len(results)} 条记录")
            return results
        except Exception as e:
            self.logger.error(f"查询执行失败: {str(e)}", exc_info=True)
            return []

    def execute_non_query(self, query: str, params: Optional[Tuple[Any, ...]] = None) -> int:
        """
        执行非查询语句（INSERT, UPDATE, DELETE 等）

        Args:
            query: SQL 语句
            params: 查询参数（用于参数化查询）

        Returns:
            int: 受影响的行数，失败返回 -1
        """
        if not self._ensure_connection():
            self.logger.error("无法执行命令，数据库连接失败")
            return -1

        try:
            if params:
                self.cursor.execute(query, params)
            else:
                self.cursor.execute(query)
            
            self.connection.commit()
            affected_rows = self.cursor.rowcount
            self.logger.info(f"命令执行成功，影响 {affected_rows} 行")
            return affected_rows
        except Exception as e:
            self.connection.rollback()
            self.logger.error(f"命令执行失败: {str(e)}", exc_info=True)
            return -1

    def __enter__(self):
        """上下文管理器入口，用于 with 语句"""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口，用于 with 语句"""
        self.close()
        if exc_type:
            self.logger.error(f"数据库操作发生异常: {exc_val}", exc_info=(exc_type, exc_val, exc_tb))
        return False

    def insert(self, table_name: str, data: Dict[str, Any]) -> int:
        """
        插入单行数据到指定表

        Args:
            table_name: 目标表名
            data: 字典形式的列名-值对

        Returns:
            int: 受影响的行数，失败返回 -1
        """
        if not data:
            self.logger.warning("插入数据为空，不执行操作")
            return 0

        columns = [f"[{col}]" for col in data.keys()]
        placeholders = ["?"] * len(data)
        sql = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({', '.join(placeholders)})"
        params = tuple(data.values())
        
        return self.execute_non_query(sql, params)

    def batch_insert(self, table_name: str, data_list: List[Dict[str, Any]]) -> int:
        """
        批量插入多行数据到指定表

        Args:
            table_name: 目标表名
            data_list: 字典列表，每个字典为一行数据

        Returns:
            int: 受影响的行数，失败返回 -1
        """
        if not data_list:
            self.logger.warning("批量插入数据为空，不执行操作")
            return 0

        # 确保所有字典的键相同
        columns = [f"[{col}]" for col in data_list[0].keys()]
        placeholders = ["?"] * len(columns)
        sql = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({', '.join(placeholders)})"
        
        # 提取所有值的元组列表
        params_list = [tuple(row.values()) for row in data_list]
        
        try:
            if not self._ensure_connection():
                self.logger.error("无法执行批量插入，数据库连接失败")
                return -1
                
            self.cursor.executemany(sql, params_list)
            self.connection.commit()
            affected_rows = self.cursor.rowcount
            self.logger.info(f"批量插入成功，影响 {affected_rows} 行")
            return affected_rows
        except Exception as e:
            self.connection.rollback()
            self.logger.error(f"批量插入失败: {str(e)}", exc_info=True)
            return -1

    def update(self, table_name: str, set_data: Dict[str, Any], where_clause: str, where_params: Optional[Tuple[Any, ...]] = None) -> int:
        """
        更新指定表中的数据

        Args:
            table_name: 目标表名
            set_data: 字典形式的列名-值对，表示要更新的字段和值
            where_clause: WHERE 子句，指定更新条件
            where_params: WHERE 子句的参数

        Returns:
            int: 受影响的行数，失败返回 -1
        """
        if not set_data:
            self.logger.warning("更新数据为空，不执行操作")
            return 0

        set_clause = ", ".join([f"[{col}] = ?" for col in set_data.keys()])
        sql = f"UPDATE {table_name} SET {set_clause} WHERE {where_clause}"
        
        # 合并 SET 和 WHERE 的参数
        params = tuple(set_data.values())
        if where_params:
            params += where_params
        
        return self.execute_non_query(sql, params)

    def delete(self, table_name: str, where_clause: str, where_params: Optional[Tuple[Any, ...]] = None) -> int:
        """
        从指定表中删除数据

        Args:
            table_name: 目标表名
            where_clause: WHERE 子句，指定删除条件
            where_params: WHERE 子句的参数

        Returns:
            int: 受影响的行数，失败返回 -1
        """
        sql = f"DELETE FROM {table_name} WHERE {where_clause}"
        return self.execute_non_query(sql, where_params)

    def execute_stored_procedure(self, proc_name: str, params: Optional[Tuple[Any, ...]] = None) -> List[Dict[str, Any]]:
        """
        执行存储过程

        Args:
            proc_name: 存储过程名称
            params: 存储过程参数

        Returns:
            List[Dict[str, Any]]: 存储过程返回的结果集
        """
        if not self._ensure_connection():
            self.logger.error("无法执行存储过程，数据库连接失败")
            return []

        try:
            sql = f"EXEC {proc_name}"
            if params:
                sql += " " + ", ".join(["?"] * len(params))
                self.cursor.execute(sql, params)
            else:
                self.cursor.execute(sql)
                
            results = []
            # 处理多个结果集
            while True:
                columns = [column[0] for column in self.cursor.description] if self.cursor.description else []
                if columns:
                    results.append([dict(zip(columns, row)) for row in self.cursor.fetchall()])
                else:
                    results.append({"rowcount": self.cursor.rowcount})
                    
                if not self.cursor.nextset():
                    break
                    
            self.logger.info(f"存储过程 {proc_name} 执行成功")
            return results
        except Exception as e:
            self.logger.error(f"存储过程执行失败: {str(e)}", exc_info=True)
            return []

    def begin_transaction(self) -> None:
        """
        开始事务
        将自动提交设置为False，需要手动调用commit()或rollback()
        """
        if self._ensure_connection():
            self.connection.autocommit = False
            self.logger.info("事务已开始")

    def commit(self) -> None:
        """
        提交事务
        """
        if self.connection:
            self.connection.commit()
            self.logger.info("事务已提交")
            self.connection.autocommit = True

    def rollback(self) -> None:
        """
        回滚事务
        """
        if self.connection:
            self.connection.rollback()
            self.logger.info("事务已回滚")
            self.connection.autocommit = True

    def get_table_schema(self, table_name: str) -> List[Dict[str, Any]]:
        """
        获取表结构信息

        Args:
            table_name: 表名

        Returns:
            List[Dict[str, Any]]: 包含列信息的字典列表
        """
        query = """
        SELECT COLUMN_NAME, DATA_TYPE, IS_NULLABLE, COLUMN_DEFAULT, CHARACTER_MAXIMUM_LENGTH
        FROM INFORMATION_SCHEMA.COLUMNS
        WHERE TABLE_NAME = ?
        ORDER BY ORDINAL_POSITION
        """
        return self.execute_query(query, (table_name,))

    def is_stored_procedure_exists(self, proc_name: str) -> bool:
        """
        检查存储过程是否存在

        Args:
            proc_name: 存储过程名称

        Returns:
            bool: 存储过程存在返回True，否则返回False
        """
        query = """
        SELECT COUNT(*) AS ProcCount
        FROM INFORMATION_SCHEMA.ROUTINES
        WHERE ROUTINE_NAME = ? AND ROUTINE_TYPE = 'PROCEDURE'
        """
        result = self.execute_query(query, (proc_name,))
        return result[0]['ProcCount'] > 0 if result else False

    def is_column_exists(self, table_name: str, column_name: str) -> bool:
        """
        检查列是否存在于表中

        Args:
            table_name: 表名
            column_name: 列名

        Returns:
            bool: 列存在返回True，否则返回False
        """
        query = """
        SELECT COUNT(*) AS ColumnCount
        FROM INFORMATION_SCHEMA.COLUMNS
        WHERE TABLE_NAME = ? AND COLUMN_NAME = ?
        """
        result = self.execute_query(query, (table_name, column_name))
        return result[0]['ColumnCount'] > 0 if result else False

    def bulk_upsert(self, table_name: str, data_list: List[Dict[str, Any]], merge_condition: str, batch_size: int = 1000) -> int:
        """
        批量 upsert 数据（插入或更新）

        Args:
            table_name: 目标表名
            data_list: 字典列表，每个字典为一行数据
            merge_condition: MERGE 语句的 ON 条件，例如 "target.ID = source.ID"
            batch_size: 每批处理的记录数

        Returns:
            int: 受影响的行数，失败返回-1
        """
        if not data_list:
            self.logger.warning("批量upsert数据为空，不执行操作")
            return 0

        columns = [f"[{col}]" for col in data_list[0].keys()]
        placeholders = [f"source.[{col}]" for col in data_list[0].keys()]
        set_clause = ", ".join([f"target.[{col}] = source.[{col}]" for col in data_list[0].keys() if col not in merge_condition])

        sql = f"""
        MERGE INTO {table_name} AS target
        USING (VALUES ({', '.join(['?'] * len(columns))})) AS source ({', '.join(columns)})
        ON {merge_condition}
        WHEN MATCHED THEN
            UPDATE SET {set_clause}
        WHEN NOT MATCHED THEN
            INSERT ({', '.join(columns)})
            VALUES ({', '.join(placeholders)});
        """

        total_affected = 0
        try:
            if not self._ensure_connection():
                self.logger.error("无法执行批量upsert，数据库连接失败")
                return -1

            for i in range(0, len(data_list), batch_size):
                batch = data_list[i:i+batch_size]
                params_list = [tuple(row.values()) for row in batch]
                self.cursor.executemany(sql, params_list)
                self.connection.commit()
                affected = self.cursor.rowcount
                total_affected += affected
                self.logger.info(f"已upsert {affected} 条记录，累计 {total_affected} 条")

            self.logger.info(f"批量upsert完成，共影响 {total_affected} 行")
            return total_affected
        except Exception as e:
            self.connection.rollback()
            self.logger.error(f"批量upsert失败: {str(e)}", exc_info=True)
            return -1

    def create_index(self, table_name: str, index_name: str, columns: List[str], is_unique: bool = False) -> bool:
        """
        创建索引

        Args:
            table_name: 表名
            index_name: 索引名称
            columns: 索引列列表
            is_unique: 是否唯一索引

        Returns:
            bool: 创建成功返回True
        """
        if not columns:
            self.logger.warning("索引列列表不能为空")
            return False

        unique_clause = "UNIQUE" if is_unique else ""
        columns_str = ", ".join([f"[{col}]" for col in columns])
        sql = f"CREATE {unique_clause} INDEX {index_name} ON {table_name} ({columns_str})"
        return self.execute_non_query(sql) > 0

    def drop_index(self, table_name: str, index_name: str) -> bool:
        """
        删除索引

        Args:
            table_name: 表名
            index_name: 索引名称

        Returns:
            bool: 删除成功返回True
        """
        sql = f"DROP INDEX {index_name} ON {table_name}"
        return self.execute_non_query(sql) > 0

    def bulk_insert(self, table_name: str, data_list: List[Dict[str, Any]], batch_size: int = 1000) -> int:
        """
        批量插入数据，支持分批处理大数据集

        Args:
            table_name: 目标表名
            data_list: 字典列表，每个字典为一行数据
            batch_size: 每批插入的记录数

        Returns:
            int: 受影响的行数，失败返回-1
        """
        if not data_list:
            self.logger.warning("批量插入数据为空，不执行操作")
            return 0

        # 确保所有字典的键相同
        columns = [f"[{col}]" for col in data_list[0].keys()]
        placeholders = ["?"] * len(columns)
        sql = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({', '.join(placeholders)})"
        
        total_affected = 0
        try:
            if not self._ensure_connection():
                self.logger.error("无法执行批量插入，数据库连接失败")
                return -1
                
            # 分批处理
            for i in range(0, len(data_list), batch_size):
                batch = data_list[i:i+batch_size]
                params_list = [tuple(row.values()) for row in batch]
                self.cursor.executemany(sql, params_list)
                self.connection.commit()
                affected = self.cursor.rowcount
                total_affected += affected
                self.logger.info(f"已插入 {affected} 条记录，累计 {total_affected} 条")
                
            self.logger.info(f"批量插入完成，共影响 {total_affected} 行")
            return total_affected
        except Exception as e:
            self.connection.rollback()
            self.logger.error(f"批量插入失败: {str(e)}", exc_info=True)
            return -1

    def is_table_exists(self, table_name: str) -> bool:
        """
        检查表是否存在

        Args:
            table_name: 表名

        Returns:
            bool: 表存在返回True，否则返回False
        """
        query = """
        SELECT COUNT(*) AS TableCount
        FROM INFORMATION_SCHEMA.TABLES
        WHERE TABLE_NAME = ?
        """
        result = self.execute_query(query, (table_name,))
        return result[0]['TableCount'] > 0 if result else False

    def bulk_update(self, table_name: str, update_data_list: List[Tuple[Dict[str, Any], str, Tuple[Any, ...]]], batch_size: int = 1000) -> int:
        """
        批量更新数据

        Args:
            table_name: 目标表名
            update_data_list: 更新数据列表，每个元素为元组(更新数据字典, where条件, where参数)
            batch_size: 每批更新的记录数

        Returns:
            int: 受影响的总行数
        """
        if not update_data_list:
            self.logger.warning("批量更新数据为空，不执行操作")
            return 0

        total_affected = 0
        try:
            if not self._ensure_connection():
                self.logger.error("无法执行批量更新，数据库连接失败")
                return -1

            # 分批处理
            for i in range(0, len(update_data_list), batch_size):
                batch = update_data_list[i:i+batch_size]
                self.begin_transaction()
                batch_affected = 0

                for set_data, where_clause, where_params in batch:
                    set_clause = ", ".join([f"[{col}] = ?" for col in set_data.keys()])
                    sql = f"UPDATE {table_name} SET {set_clause} WHERE {where_clause}"
                    
                    params = tuple(set_data.values())
                    if where_params:
                        params += where_params
                    
                    affected = self.execute_non_query(sql, params)
                    if affected > 0:
                        batch_affected += affected

                self.commit()
                total_affected += batch_affected
                self.logger.info(f"已更新 {batch_affected} 条记录，累计 {total_affected} 条")

            self.logger.info(f"批量更新完成，共影响 {total_affected} 行")
            return total_affected
        except Exception as e:
            self.rollback()
            self.logger.error(f"批量更新失败: {str(e)}", exc_info=True)
            return -1

    def bulk_delete(self, table_name: str, where_clauses: List[Tuple[str, Tuple[Any, ...]]], batch_size: int = 1000) -> int:
        """
        批量删除数据

        Args:
            table_name: 目标表名
            where_clauses: 删除条件列表，每个元素为元组(where条件, where参数)
            batch_size: 每批删除的记录数

        Returns:
            int: 受影响的总行数
        """
        if not where_clauses:
            self.logger.warning("删除条件列表为空，不执行操作")
            return 0

        total_affected = 0
        try:
            if not self._ensure_connection():
                self.logger.error("无法执行批量删除，数据库连接失败")
                return -1

            # 分批处理
            for i in range(0, len(where_clauses), batch_size):
                batch = where_clauses[i:i+batch_size]
                self.begin_transaction()
                batch_affected = 0

                for where_clause, where_params in batch:
                    sql = f"DELETE FROM {table_name} WHERE {where_clause}"
                    affected = self.execute_non_query(sql, where_params)
                    if affected > 0:
                        batch_affected += affected

                self.commit()
                total_affected += batch_affected
                self.logger.info(f"已删除 {batch_affected} 条记录，累计 {total_affected} 条")

            self.logger.info(f"批量删除完成，共影响 {total_affected} 行")
            return total_affected
        except Exception as e:
            self.rollback()
            self.logger.error(f"批量删除失败: {str(e)}", exc_info=True)
            return -1

    def get_indexes(self, table_name: str) -> List[Dict[str, Any]]:
        """
        获取表的索引信息

        Args:
            table_name: 表名

        Returns:
            List[Dict[str, Any]]: 索引信息列表
        """
        query = """
        SELECT
            i.name AS index_name,
            COL_NAME(ic.object_id, ic.column_id) AS column_name,
            i.is_unique,
            i.type_desc
        FROM
            sys.indexes i
        INNER JOIN
            sys.index_columns ic ON i.object_id = ic.object_id AND i.index_id = ic.index_id
        WHERE
            OBJECT_NAME(i.object_id) = ?
        ORDER BY
            i.name, ic.key_ordinal
        """
        return self.execute_query(query, (table_name,))